# trace generated using paraview version 5.10.0
#
# import the simple module from the paraview
from paraview.simple import *
import concurrent.futures
import os

def show_one_SR(working_dir, id):
    # create a new 'XML Structured Grid Reader'
    target_vts = XMLStructuredGridReader(
        FileName=[working_dir+"/"+id+".vts"])

    # get active view
    renderView1 = GetActiveViewOrCreate('RenderView')

    # show data in view
    target_vtsDisplay = Show(target_vts, renderView1)

    # trace defaults for the display properties.
    target_vtsDisplay.Representation = 'Outline'

    # reset view to fit data
    renderView1.ResetCamera()

    # update the view to ensure updated data information
    renderView1.Update()

    # hide data in view
    Hide(target_vts, renderView1)

    # create a new 'Iso Volume'
    isoVolume1 = IsoVolume(Input=target_vts)

    # Properties modified on isoVolume1
    isoVolume1.InputScalars = ['POINTS', 'define_marker_rho']
    isoVolume1.ThresholdRange = [0.0, 50.0]

    # show data in view
    isoVolume1Display = Show(isoVolume1, renderView1)

    # trace defaults for the display properties.
    isoVolume1Display.Representation = 'Surface'

    # hide data in view
    Hide(isoVolume1, renderView1)

    # create a new 'Iso Volume'
    isoVolume2 = IsoVolume(Input=isoVolume1)

    # Properties modified on isoVolume2
    isoVolume2.InputScalars = ['POINTS', 'define_marker_theta']
    isoVolume2.ThresholdRange = [0.0, 50.0]

    # show data in view
    isoVolume2Display = Show(isoVolume2, renderView1)

    # trace defaults for the display properties.
    isoVolume2Display.Representation = 'Surface'

    # hide data in view
    Hide(isoVolume2, renderView1)

    # create a new 'Contour'
    contour1 = Contour(Input=isoVolume2)

    # Properties modified on contour1
    contour1.Isosurfaces = [-0.002]
    contour1.ComputeNormals = 1
    contour1.GenerateTriangles = 0

    # show data in view
    contour1Display = Show(contour1, renderView1)

    # trace defaults for the display properties.
    contour1Display.Representation = 'Surface'

    # turn off scalar coloring
    ColorBy(contour1Display, None)

    # update the view to ensure updated data information
    renderView1.Update()

    # change solid color
    contour1Display.DiffuseColor = [1.0, 0.0, 0.0]

    # set active source
    SetActiveSource(isoVolume2)

    # create a new 'Iso Volume'
    isoVolume3 = IsoVolume(Input=isoVolume2)

    # Properties modified on isoVolume3
    isoVolume3.ThresholdRange = [0.0, 50.0]

    # show data in view
    isoVolume3Display = Show(isoVolume3, renderView1)

    # trace defaults for the display properties.
    isoVolume3Display.Representation = 'Surface'

    # hide data in view
    Hide(isoVolume1, renderView1)

    # show color bar/color legend
    isoVolume3Display.SetScalarBarVisibility(renderView1, False)

    # update the view to ensure updated data information
    renderView1.Update()

    # turn off scalar coloring
    ColorBy(isoVolume3Display, None)

    # reset view to fit data
    renderView1.ResetCamera()

    # current camera placement for renderView1
    # renderView1.CameraPosition = [-1.1733679081156507, -2.721868099824966, -3.2596504885633077]
    renderView1.CameraPosition = [-2, -2, -2]
    # renderView1.CameraFocalPoint = [0.7095798552036285, 0.21344321966171265, 1.5515306293964386]
    # renderView1.CameraViewUp = [-0.5334919059913795, 0.7986507272929785, -0.2784661595887045]
    # renderView1.CameraParallelScale = 0.1
    renderView1.CameraViewUp = [0.0, 1.0, 0.0]

    renderView1.ResetCamera()

    # Properties modified on renderView1
    renderView1.OrientationAxesVisibility = 0

    # save screenshot
    SaveScreenshot(working_dir+"/"+id+"SR.png", renderView1, ImageResolution=[2000, 2000],
                    TransparentBackground=1)
    print("finish SR visualization for "+id)

    # destroy isoVolume1

    Delete(isoVolume3)
    del isoVolume3

    Delete(isoVolume2)
    del isoVolume2

    Delete(contour1)
    del contour1

    Delete(isoVolume1)
    del isoVolume1

    Delete(target_vts)
    del target_vts


def show_all_SR(working_dir):
    # Get all files in the directory
    files = os.listdir(working_dir)

    # Filter files with .vts extension and get only the file names
    vts_file_names = [os.path.splitext(file)[0]
                      for file in files if file.endswith('.vts')]

    # Create a ProcessPoolExecutor
    with concurrent.futures.ProcessPoolExecutor(max_workers=3) as executor:
        # Submit tasks to the executor
        futures = [executor.submit(show_one_SR, working_dir, id) for id in vts_file_names]

    # Wait for all tasks to complete
    concurrent.futures.wait(futures)
        
if __name__ == '__main__':
    from DEAD.AutoDecoder.Config import result_save_path
    show_all_SR(f"{result_save_path}/VTK")
